{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0e81b124",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/retrievers/vectara_auto_retriever.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "307804a3-c02b-4a57-ac0d-172c30ddc851",
   "metadata": {},
   "source": [
    "# Auto-Retrieval from a Vectara Index\n",
    "\n",
    "This guide shows how to perform **auto-retrieval** in LlamaIndex with Vectara. \n",
    "\n",
    "Given a natural language query, we first use the LLM to infer a set of metadata filters as well as the right query string to pass to the Vectara Index.\n",
    "\n",
    "This allows for more dynamic, expressive forms of retrieval beyond top-k semantic search. The relevant context for a given query may only require filtering on a metadata tag, or require a joint combination of filtering + semantic search within the filtered set, or just raw semantic search."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "## Setup "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31faecfb",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d3d6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d48af8e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdf284a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.schema import TextNode\n",
    "from llama_index.indices.managed.types import ManagedIndexQueryMode\n",
    "from llama_index.indices import VectaraIndex\n",
    "from llama_index.indices.managed.vectara import VectaraAutoRetriever\n",
    "\n",
    "from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n",
    "\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index import ServiceContext"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41aa106b-8261-4a01-97c6-1b037dffa1b4",
   "metadata": {},
   "source": [
    "## Defining Some Sample Data\n",
    "\n",
    "We insert some sample nodes containing text chunks into Vectara. Note that each `TextNode` not only contains the text, but also certain metadata fields like year, director, rating or genre. \n",
    "In Vectara you will need to [define](https://docs.vectara.com/docs/learn/metadata-search-filtering/filter-overview) these metadata fields in your coprus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = [\n",
    "    TextNode(\n",
    "        text=(\n",
    "            \"A pragmatic paleontologist touring an almost complete theme park on an island \"\n",
    "            + \"in Central America is tasked with protecting a couple of kids after a power \"\n",
    "            + \"failure causes the park's cloned dinosaurs to run loose.\"\n",
    "        ),\n",
    "        metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": \"science fiction\"},\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=(\n",
    "            \"A thief who steals corporate secrets through the use of dream-sharing technology \"\n",
    "            + \"is given the inverse task of planting an idea into the mind of a C.E.O., \"\n",
    "            + \"but his tragic past may doom the project and his team to disaster.\"\n",
    "        ),\n",
    "        metadata={\n",
    "            \"year\": 2010,\n",
    "            \"director\": \"Christopher Nolan\",\n",
    "            \"rating\": 8.2,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"Barbie suffers a crisis that leads her to question her world and her existence.\",\n",
    "        metadata={\n",
    "            \"year\": 2023,\n",
    "            \"director\": \"Greta Gerwig\",\n",
    "            \"genre\": \"fantasy\",\n",
    "            \"rating\": 9.5,\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=(\n",
    "            \"A cowboy doll is profoundly threatened and jealous when a new spaceman action \"\n",
    "            + \"figure supplants him as top toy in a boy's bedroom.\"\n",
    "        ),\n",
    "        metadata={\"year\": 1995, \"genre\": \"animated\", \"rating\": 8.3},\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=(\n",
    "            \"When Woody is stolen by a toy collector, Buzz and his friends set out on a \"\n",
    "            + \"rescue mission to save Woody before he becomes a museum toy property with his \"\n",
    "            + \"roundup gang Jessie, Prospector, and Bullseye. \"\n",
    "        ),\n",
    "        metadata={\"year\": 1999, \"genre\": \"animated\", \"rating\": 7.9},\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=(\n",
    "            \"The toys are mistakenly delivered to a day-care center instead of the attic \"\n",
    "            + \"right before Andy leaves for college, and it's up to Woody to convince the \"\n",
    "            + \"other toys that they weren't abandoned and to return home.\"\n",
    "        ),\n",
    "        metadata={\"year\": 2010, \"genre\": \"animated\", \"rating\": 8.3},\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8bd70be-57c7-49e2-990b-ad9a876710fb",
   "metadata": {},
   "source": [
    "## Build Vectara Managed Index\n",
    "\n",
    "Now we load our sample data into the Vectara Index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ecb26f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM is explicitly disabled. Using MockLLM.\n"
     ]
    }
   ],
   "source": [
    "index = VectaraIndex(nodes=nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acab9d9c",
   "metadata": {},
   "source": [
    "## Setup OpenAI\n",
    "\n",
    "Auto-retrieval uses an LLM to convert the natural language query into a shorter query and meta data filtering conditions. We will be using the OpenAI LLM, so let's set that up here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcac4eda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import openai\n",
    "\n",
    "if not os.environ.get(\"OPENAI_API_KEY\", None):\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n",
    "\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c793dc45-5087-4dcb-b0d3-85b8e718539f",
   "metadata": {},
   "source": [
    "## Define `VectorStoreInfo`\n",
    "\n",
    "We define `VectorStoreInfo`, which contains a structured description of the metadata filters suported. This information is later on usedin the auto-retrieval prompt where the LLM infers metadata filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bedbb693-725f-478f-be26-fa7180ea38b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store_info = VectorStoreInfo(\n",
    "    content_info=\"information about a movie\",\n",
    "    metadata_info=[\n",
    "        MetadataInfo(\n",
    "            name=\"genre\",\n",
    "            description=\"The genre of the movie. One of ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
    "            type=\"string\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"year\",\n",
    "            description=\"The year the movie was released\",\n",
    "            type=\"integer\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"director\",\n",
    "            description=\"The name of the movie director\",\n",
    "            type=\"string\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"rating\",\n",
    "            description=\"A 1-10 rating for the movie\",\n",
    "            type=\"float\",\n",
    "        ),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32808a60-7bab-4e9e-944c-cfe2ed0b0e2e",
   "metadata": {},
   "source": [
    "## Running over some sample data\n",
    "\n",
    "Now let's create a `VectaraAutoRetriever` instance and run some example queries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1028bc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "service_context = ServiceContext.from_defaults(\n",
    "    llm=OpenAI(model=\"gpt-4-1106-preview\", temperature=0)\n",
    ")\n",
    "retriever = VectaraAutoRetriever(\n",
    "    index,\n",
    "    vector_store_info=vector_store_info,\n",
    "    service_context=service_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb18e9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[NodeWithScore(node=TextNode(id_='19ce35022e05e1d8672d1c609ce71b7b0958e1ecfc61717f54d3df26bf43b2392dc12b4148842e582d50ba813aa9381141e9ddefd4f6315da36e85efb77844f5', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '79', 'year': '2023', 'director': 'Greta Gerwig', 'genre': 'fantasy', 'rating': '9.5'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='ad0e4a83d372f73b0226cb81a9ccdbd4c77526a6a70cb2d032bd5e9331441c6d', text='Barbie suffers a crisis that leads her to question her world and her existence.', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.59363246)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.retrieve(\"movie directed by Greta Gerwig\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95d4a0e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[NodeWithScore(node=TextNode(id_='809f292b591d54338f7d8cbb3db72af48bb8a0b9de0d10cb90ab1e340b9670886387816d62acc6a61cf2a65938acfb7eb95cc8d4a65d54c72b9bd10b489fa9ee', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '209', 'year': '2010', 'genre': 'animated', 'rating': '8.3'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='c1c9f42ca5857db9d154c05d68139a5d8515b1c90e6032d41deb8fdc2138ff40', text=\"The toys are mistakenly delivered to a day-care center instead of the attic right before Andy leaves for college, and it's up to Woody to convince the other toys that they weren't abandoned and to return home.\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.740946),\n",
       " NodeWithScore(node=TextNode(id_='4bc116ca95b0802048d37b5807ff8b177b7ae0617207d6f56b998302b8c400cc99deedee8bdb1ccfb80107d3f0b20343c4be59ff0ec5c60312dc0ee5338d9f17', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '129', 'year': '1995', 'genre': 'animated', 'rating': '8.3'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='35a858f68a9181d0a97045cb4dc92db5a02f307686ce20405e3545d1a0cf5020', text=\"A cowboy doll is profoundly threatened and jealous when a new spaceman action figure supplants him as top toy in a boy's bedroom.\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.67997515)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.retrieve(\"movie about toys with a rating above 8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a3eb674",
   "metadata": {},
   "source": [
    "We can also include standard `VectaraRetriever` arguments in the `VectaraAutoRetriever`. For example, if we want to include a `filter` that would be added to any additional filtering from the query itself, we can do it as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ed993b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[NodeWithScore(node=TextNode(id_='809f292b591d54338f7d8cbb3db72af48bb8a0b9de0d10cb90ab1e340b9670886387816d62acc6a61cf2a65938acfb7eb95cc8d4a65d54c72b9bd10b489fa9ee', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '209', 'year': '2010', 'genre': 'animated', 'rating': '8.3'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='c1c9f42ca5857db9d154c05d68139a5d8515b1c90e6032d41deb8fdc2138ff40', text=\"The toys are mistakenly delivered to a day-care center instead of the attic right before Andy leaves for college, and it's up to Woody to convince the other toys that they weren't abandoned and to return home.\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.740946),\n",
       " NodeWithScore(node=TextNode(id_='4bc116ca95b0802048d37b5807ff8b177b7ae0617207d6f56b998302b8c400cc99deedee8bdb1ccfb80107d3f0b20343c4be59ff0ec5c60312dc0ee5338d9f17', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '129', 'year': '1995', 'genre': 'animated', 'rating': '8.3'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='35a858f68a9181d0a97045cb4dc92db5a02f307686ce20405e3545d1a0cf5020', text=\"A cowboy doll is profoundly threatened and jealous when a new spaceman action figure supplants him as top toy in a boy's bedroom.\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.67997515)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever = VectaraAutoRetriever(\n",
    "    index,\n",
    "    vector_store_info=vector_store_info,\n",
    "    service_context=service_context,\n",
    "    filter=\"doc.rating > 8\",\n",
    ")\n",
    "retriever.retrieve(\"movie about toys\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11c2f0a2",
   "metadata": {},
   "source": [
    "Now let's try with MMR (max marginal relevance). To demonstrate the maximum effect we will use mmr_diversity_bias value of 1.0 (maximum diversity), noting that typical value is usually 0.2 or 0.3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d62ff971",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[NodeWithScore(node=TextNode(id_='809f292b591d54338f7d8cbb3db72af48bb8a0b9de0d10cb90ab1e340b9670886387816d62acc6a61cf2a65938acfb7eb95cc8d4a65d54c72b9bd10b489fa9ee', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '209', 'year': '2010', 'genre': 'animated', 'rating': '8.3'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='c1c9f42ca5857db9d154c05d68139a5d8515b1c90e6032d41deb8fdc2138ff40', text=\"The toys are mistakenly delivered to a day-care center instead of the attic right before Andy leaves for college, and it's up to Woody to convince the other toys that they weren't abandoned and to return home.\", start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.740946),\n",
       " NodeWithScore(node=TextNode(id_='2bc612264c3ccf53264b02765eb477d6d4aaeca6aa6261e42611307ddd8079a340f291436bd3f7b2e84a940ee4d4c51333f6d52e0c7459096501aff6a62deef8', embedding=None, metadata={'lang': 'eng', 'offset': '0', 'len': '220', 'year': '2010', 'director': 'Christopher Nolan', 'rating': '8.2'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='a7925bf91469af831471913fd164a40c6711a6e6b236827b1a7c971e4b2cfb5e', text='A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O., but his tragic past may doom the project and his team to disaster.', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=-0.7325192)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever = VectaraAutoRetriever(\n",
    "    index,\n",
    "    vector_store_info=vector_store_info,\n",
    "    service_context=service_context,\n",
    "    filter=\"doc.rating > 8\",\n",
    "    vectara_query_mode=\"mmr\",\n",
    "    mmr_k=50,\n",
    "    mmr_diversity_bias=1,\n",
    ")\n",
    "retriever.retrieve(\"movie about toys\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85873347",
   "metadata": {},
   "source": [
    "We can see that the results are reranked with MMR to create more diversity, and instead of two \"toy story\" results we get a first result about toy story and another one for the movie inception."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "0ac390d292208ca2380c85f5bce7ded36a7a25670a97c40b8009630eb36cb06e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
